# @Version        : 1.0
# @Update Time    : 2024/12/30 23:47
# @File           : query_builder.py
# @IDE            : PyCharm
# @Desc           : 查询构建器，用于构建SQLAlchemy查询条件，基于advanced-alchemy的filters模块

from datetime import datetime
from enum import Enum
from typing import List, Any, Union, Optional, Dict, TypeVar, Generic, Set, Type
from sqlalchemy import and_, or_, ColumnElement, select
from sqlalchemy.orm import InstrumentedAttribute
from advanced_alchemy import filters
from src.log import logger
from src.common.typing.common_typing import ORMModel
from src.core.exception import BaseAppException


class QueryType(Enum):
    """
    查询类型枚举

    支持的查询类型：
    - LIKE: 模糊匹配
    - EQ: 等于
    - NE: 不等于
    - GE: 大于等于
    - LT: 小于
    - LE: 小于等于
    - IN: 包含
    - NOT_IN: 不包含
    - BETWEEN: 范围查询
    - ILIKE: 不区分大小写的模糊匹配
    - STARTSWITH: 开头匹配
    - ENDSWITH: 结尾匹配
    - IENDSWITH: 不区分大小写的结尾匹配
    - DATEEQ: 日期相等
    - SEARCH: 全文搜索
    - NOT_EXISTS: 不存在
    - EXISTS: 存在
    - BEFORE_AFTER: 日期范围查询（不包含边界）
    - ON_BEFORE_AFTER: 日期范围查询（包含边界）
    """

    LIKE = "like"
    EQ = "eq"
    NE = "ne"
    GE = "ge"
    LT = "lt"
    LE = "le"
    IN = "in"
    NOT_IN = "notin"
    BETWEEN = "between"
    ILIKE = "ilike"
    STARTSWITH = "startswith"
    ENDSWITH = "endswith"
    IENDSWITH = "iendswith"
    DATEEQ = "dateeq"
    SEARCH = "search"
    NOT_EXISTS = "not_exists"
    EXISTS = "exists"
    BEFORE_AFTER = "before_after"
    ON_BEFORE_AFTER = "on_before_after"


T = TypeVar("T")


class FilterCriteria:
    """
    查询条件类

    用于封装查询条件信息，包括：
    - 查询类型
    - 查询值
    - 字段别名
    - 模型名称
    """

    def __init__(
        self,
        query_type: QueryType,
        value: Any,
        alias: Optional[str] = None,
        model_name: Optional[str] = None,
        ignore_case: bool = False,
    ):
        """
        初始化查询条件

        Args:
            query_type: 查询类型
            value: 查询值
            alias: 字段别名
            model_name: 模型名称
            ignore_case: 是否忽略大小写
        """
        self.query_type = query_type
        self.value = value
        self.alias = alias
        self.model_name = model_name
        self.ignore_case = ignore_case

    def to_dict(self) -> Dict[str, Any]:
        """
        转换为字典格式

        Returns:
            Dict[str, Any]: 字典格式的查询条件
        """
        return {
            "query_type": self.query_type.value,
            "value": self.value,
            "alias": self.alias,
            "model_name": self.model_name,
            "ignore_case": self.ignore_case,
        }


class QueryBuilder:
    """
    查询构建器

    用于构建SQLAlchemy查询条件，支持：
    - 单条件查询
    - 多条件组合查询
    - 日期范围查询（包含和不包含边界）
    - 多模型关联查询
    - 全文搜索
    - 存在性查询
    """

    def __init__(self, model: ORMModel, tables: Optional[Dict[str, ORMModel]] = None):
        """
        初始化查询构建器

        Args:
            model: SQLAlchemy模型类或别名对象
            tables: 多模型查询时的模型字典
        """
        self.model = model
        self.tables = tables or {}
        self.conditions: List[Union[ColumnElement, filters.StatementFilter]] = []
        self.order_by: List[filters.OrderBy] = []
        self.limit_offset: Optional[filters.LimitOffset] = None

    def add_condition(
        self,
        field: str,
        query_type: Union[QueryType, str],
        value: Any,
        field_alias: Optional[str] = None,
        model_name: Optional[str] = None,
        ignore_case: bool = False,
    ) -> "QueryBuilder":
        """
        添加查询条件

        Args:
            field: 字段名
            query_type: 查询类型
            value: 查询值
            field_alias: 字段别名
            model_name: 模型名称
            ignore_case: 是否忽略大小写

        Returns:
            QueryBuilder: 支持链式调用

        Raises:
            BaseAppException: 查询类型无效或模型不存在时抛出
        """
        # 处理空值
        if isinstance(value, str):
            value = value.strip()
        if value is None or value == "":
            return self

        # 获取字段属性
        field = field_alias or field
        try:
            if model_name:
                if model_name not in self.tables:
                    raise BaseAppException(msg=f"模型 '{model_name}' 不存在")
                attr: InstrumentedAttribute = getattr(self.tables[model_name], field)
            else:
                attr: InstrumentedAttribute = getattr(self.model, field, None)
                if not attr:
                    logger.warning(f"字段 '{field}' 在模型中不存在")
                    return self
        except AttributeError as e:
            logger.error(f"获取字段属性失败: {str(e)}")
            return self

        # 验证查询类型
        if isinstance(query_type, QueryType):
            query_type = query_type.value
        filter_list = ["before_after", "on_before_after"]
        filter_list.extend(filters.VALID_OPERATORS)
        if query_type not in filter_list:
            raise BaseAppException(msg=f"不支持的查询类型: {query_type}")

        # 根据查询类型创建不同的过滤器
        if query_type == QueryType.SEARCH.value:
            self.conditions.append(filters.SearchFilter(field, value, ignore_case))
        elif query_type == QueryType.EXISTS.value:
            self.conditions.append(filters.ExistsFilter(value))
        elif query_type == QueryType.NOT_EXISTS.value:
            self.conditions.append(filters.NotExistsFilter(value))
        elif query_type == QueryType.BEFORE_AFTER.value:
            if not isinstance(value, (tuple, list)) or len(value) != 2:
                raise BaseAppException(
                    msg="日期范围查询值必须是包含两个元素的元组或列表"
                )
            before, after = value
            self.conditions.append(filters.BeforeAfter(field, before, after))
        elif query_type == QueryType.ON_BEFORE_AFTER.value:
            if not isinstance(value, (tuple, list)) or len(value) != 2:
                raise BaseAppException(
                    msg="日期范围查询值必须是包含两个元素的元组或列表"
                )
            before, after = value
            self.conditions.append(filters.OnBeforeAfter(field, before, after))
        else:
            self.conditions.append(filters.ComparisonFilter(field, query_type, value))

        return self

    def add_order_by(
        self,
        field: str,
        sort_order: str = "asc",
        model_name: Optional[str] = None,
    ) -> "QueryBuilder":
        """
        添加排序条件

        Args:
            field: 字段名
            sort_order: 排序方式，asc或desc
            model_name: 模型名称

        Returns:
            QueryBuilder: 支持链式调用
        """
        self.order_by.append(filters.OrderBy(field, sort_order))
        return self

    def add_pagination(
        self,
        limit: int,
        offset: int = 0,
    ) -> "QueryBuilder":
        """
        添加分页条件

        Args:
            limit: 每页数量
            offset: 偏移量

        Returns:
            QueryBuilder: 支持链式调用
        """
        self.limit_offset = filters.LimitOffset(limit, offset)
        return self

    def build(
        self, params: Dict[str, Any]
    ) -> List[Union[ColumnElement, filters.StatementFilter]]:
        """
        构建查询条件

        Args:
            params: 查询参数字典

        Returns:
            List[Union[ColumnElement, filters.StatementFilter]]: 查询条件列表

        Raises:
            BaseAppException: 查询条件格式无效时抛出
        """
        if not params:
            return []

        for field, query_value in params.items():
            if query_value is None or query_value == "":
                continue

            try:
                if isinstance(query_value, FilterCriteria):
                    self.add_condition(
                        field,
                        query_type=query_value.query_type,
                        value=query_value.value,
                        field_alias=query_value.alias,
                        model_name=query_value.model_name,
                        ignore_case=query_value.ignore_case,
                    )
                elif isinstance(query_value, (tuple, list)):
                    if len(query_value) < 2:
                        raise BaseAppException(msg="查询条件格式无效")
                    self.add_condition(field, *query_value)
                else:
                    self.add_condition(
                        field, query_type=QueryType.EQ.value, value=query_value
                    )
            except Exception as e:
                logger.error(f"构建查询条件失败: {str(e)}")
                continue

        return self.conditions

    def build_and(self, params: Dict[str, Any]) -> List[ColumnElement[bool]]:
        """
        构建AND组合的查询条件

        Args:
            params: 查询参数字典

        Returns:
            List[ColumnElement[bool]]: AND组合的查询条件
        """
        self.build(params)
        return [and_(*self.conditions)] if self.conditions else []

    def build_or(self, params: Dict[str, Any]) -> List[ColumnElement[bool]]:
        """
        构建OR组合的查询条件

        Args:
            params: 查询参数字典

        Returns:
            List[ColumnElement[bool]]: OR组合的查询条件
        """
        self.build(params)
        return [or_(*self.conditions)] if self.conditions else []

    def build_query(
        self,
        params: Dict[str, Any],
        *,
        order_by: Optional[List[tuple[str, str]]] = None,
        limit: Optional[int] = None,
        offset: Optional[int] = None,
    ) -> select:
        """
        构建完整的查询语句

        Args:
            params: 查询参数字典
            order_by: 排序条件列表，格式为[(field, sort_order), ...]
            limit: 每页数量
            offset: 偏移量

        Returns:
            select: SQLAlchemy查询语句
        """
        # 构建基础查询
        query = select(self.model)

        # 添加查询条件
        conditions = self.build(params)
        if conditions:
            query = query.where(and_(*conditions))

        # 添加排序
        if order_by:
            for field, sort_order in order_by:
                self.add_order_by(field, sort_order)
        for order in self.order_by:
            query = order.append_to_statement(query, self.model)

        # 添加分页
        if limit is not None:
            self.add_pagination(limit, offset or 0)
        if self.limit_offset:
            query = self.limit_offset.append_to_statement(query, self.model)

        return query
